import cv2
import numpy as np
import yaml
from pathlib import Path
import serial
import time
import threading

class CircleDetectionAndValidation:
    def __init__(self, calibration_file="hand_eye_calibration_results.txt", com_port="/dev/ttyUSB0"):
        self.camera_matrix = None
        self.dist_coeffs = None
        self.H = None  # Homography matrix
        self.load_camera_calibration()
        self.load_hand_eye_calibration(calibration_file)
        
        # 串口初始化
        self.ser = None
        self.serial_connected = self.initialize_serial(com_port)
        
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            raise Exception("Cannot open camera")
        
        # Set camera parameters
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        
        # Black circle detection parameters
        self.min_radius = 3
        self.max_radius = 5
        self.black_threshold = 70
        
        # Movement control state
        self.target_position = None  # Target position in world coordinates
        self.current_position = None  # Current position in world coordinates
        self.image_center = (325, 235)  # 图像中心坐标
        
        # 动态追踪参数
        self.tracking_enabled = False
        self.tolerance_pixels = 5  # 中心点容差像素
        self.min_movement_threshold = 0.1  # 最小移动阈值(mm)
        
        # 如果串口连接成功才初始化电机
        if self.serial_connected:
            self.initialize_motor()
        else:
            print("警告: 串口未连接，电机控制功能不可用")
    
    def load_camera_calibration(self, config_path='camera_calibration.yaml'):
        """Load camera intrinsic parameters"""
        try:
            if Path(config_path).exists():
                with open(config_path, 'r') as file:
                    calibration_data = yaml.safe_load(file)
                
                if 'camera_matrix' in calibration_data:
                    self.camera_matrix = np.array(calibration_data['camera_matrix'])
                if 'distortion_coefficients' in calibration_data:
                    self.dist_coeffs = np.array(calibration_data['distortion_coefficients'])
                
                print("Camera intrinsic parameters loaded successfully")
            else:
                print("Camera calibration file not found, using default parameters")
                self.camera_matrix = np.eye(3)
                self.dist_coeffs = np.zeros((5, 1))
                
        except Exception as e:
            print(f"Error loading camera calibration: {e}")
            self.camera_matrix = np.eye(3)
            self.dist_coeffs = np.zeros((5, 1))
    
    def initialize_serial(self, com_port):
        """初始化串口连接"""
        try:
            print(f"尝试连接串口: {com_port}...")
            self.ser = serial.Serial(
                port=com_port,
                baudrate=250000,
                timeout=1,
                write_timeout=1,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE
            )
            
            # 等待串口初始化
            time.sleep(2)
            
            # 测试串口是否真的可用
            self.ser.reset_input_buffer()
            self.ser.reset_output_buffer()
            
            print(f"串口 {com_port} 连接成功，波特率 250000")
            return True
            
        except serial.SerialException as e:
            print(f"串口连接失败: {e}")
            print("可能的原因:")
            print("1. 设备不存在: /dev/ttyUSB0")
            print("2. 权限不足，尝试: sudo chmod 666 /dev/ttyUSB0")
            print("3. 设备被其他程序占用")
            print("4. 波特率不匹配")
            self.ser = None
            return False
        except Exception as e:
            print(f"串口连接未知错误: {e}")
            self.ser = None
            return False
    
    def send_gcode_command(self, command):
        """发送G代码命令并等待返回ok"""
        if self.ser is None or not self.serial_connected:
            print(f"串口未连接，无法发送命令: {command}")
            return False
        
        try:
            # 清空输入缓冲区
            self.ser.reset_input_buffer()
            
            # 发送命令
            print(f"发送命令: {command.strip()}")
            full_command = command + '\n'
            self.ser.write(full_command.encode())
            
            # 等待响应
            start_time = time.time()
            response_lines = []
            
            while time.time() - start_time < 5:  # 5秒超时
                if self.ser.in_waiting > 0:
                    response = self.ser.readline().decode().strip()
                    if response:
                        print(f"收到响应: {response}")
                        response_lines.append(response)
                        
                        if "ok" in response.lower():
                            return True
                        elif "error" in response.lower():
                            print(f"命令错误: {response}")
                            return False
                
                time.sleep(0.1)
            
            if response_lines:
                print(f"响应超时，但收到了: {response_lines}")
            else:
                print("等待响应超时，未收到任何响应")
            return False
            
        except Exception as e:
            print(f"发送命令错误: {e}")
            return False
    
    def initialize_motor(self):
        """初始化电机，发送必要的初始化命令"""
        if not self.serial_connected:
            print("串口未连接，跳过电机初始化")
            return
        
        print("开始初始化电机...")
        
        # 先等待一下确保串口稳定
        time.sleep(1)
        
        # 发送初始化命令序列
        commands = [
            "G0 X1",      # 移动到X1位置
            "G28 X Y",    # 归零X Y轴
            "G0 X150 Y90" # 移动到初始位置
        ]
        
        for cmd in commands:
            print(f"发送初始化命令: {cmd}")
            if not self.send_gcode_command(cmd):
                print(f"电机初始化失败在命令: {cmd}")
                # 继续尝试下一个命令而不是直接返回
                continue
            time.sleep(1)  # 命令间延迟
        
        print("电机初始化完成")
    
    def move_motor_to_center(self, dx, dy):
        """控制电机移动到中心位置"""
        if not self.serial_connected:
            print("串口未连接，无法控制电机")
            return False
        
        print(f"开始移动电机: X{dx:.2f}mm, Y{dy:.2f}mm")
        
        # 设置为相对坐标模式
        if not self.send_gcode_command("G91"):
            print("设置相对坐标模式失败")
            return False
        
        time.sleep(0.5)
        
        # 发送移动命令
        move_command = f"G0 X{dx:.2f} Y{dy:.2f}"
        move_success = self.send_gcode_command(move_command)
        
        time.sleep(0.5)
        
        # 切换回绝对坐标模式
        self.send_gcode_command("G90")
        
        return move_success
    
    def load_hand_eye_calibration(self, calibration_file):
        """Load hand-eye calibration results"""
        try:
            if not Path(calibration_file).exists():
                raise FileNotFoundError(f"Hand-eye calibration file {calibration_file} not found")
            
            print(f"Loading hand-eye calibration results: {calibration_file}")
            
            with open(calibration_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # Parse homography matrix
            matrix_lines = []
            in_matrix_section = False
            
            for line in lines:
                line = line.strip()
                if '单应性矩阵' in line or 'Homography matrix' in line:
                    in_matrix_section = True
                    continue
                if in_matrix_section and line.startswith('['):
                    matrix_lines.append(line)
                if in_matrix_section and not line.startswith('[') and matrix_lines:
                    break
            
            if len(matrix_lines) == 3:
                self.H = np.array([
                    [float(x) for x in line.strip('[]').split(',')] 
                    for line in matrix_lines
                ])
                print("Homography matrix loaded successfully:")
                print(self.H)
            else:
                print("Warning: Cannot parse homography matrix, using identity matrix")
                self.H = np.eye(3)
                
        except Exception as e:
            print(f"Error loading hand-eye calibration: {e}")
            self.H = np.eye(3)
    
    def undistort_image(self, frame):
        """Undistort image"""
        if self.camera_matrix is not None and self.dist_coeffs is not None:
            try:
                h, w = frame.shape[:2]
                new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
                    self.camera_matrix, self.dist_coeffs, (w, h), 1, (w, h)
                )
                undistorted = cv2.undistort(frame, self.camera_matrix, self.dist_coeffs, None, new_camera_matrix)
                return undistorted
            except Exception as e:
                print(f"Undistortion error: {e}")
                return frame
        return frame
    
    def detect_black_circles(self, frame):
        """Detect black circles"""
        # Convert to grayscale
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur
        blurred = cv2.GaussianBlur(gray, (7, 7), 2)
        
        # Adaptive thresholding
        binary = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                     cv2.THRESH_BINARY_INV, 11, 2)
        
        # Morphological operations
        kernel = np.ones((2, 2), np.uint8)
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
        
        # Find contours
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        detected_circles = []
        
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 20 or area > 1000:
                continue
            
            perimeter = cv2.arcLength(contour, True)
            if perimeter == 0:
                continue
                
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            if circularity < 0.7:
                continue
            
            (x, y), radius = cv2.minEnclosingCircle(contour)
            center = (int(x), int(y))
            radius = int(radius)
            
            if self.min_radius <= radius <= self.max_radius:
                detected_circles.append((center[0], center[1], radius))
        
        return detected_circles, binary
    
    def pixel_to_world(self, pixel_x, pixel_y):
        """Convert pixel coordinates to world coordinates"""
        if self.H is None:
            return pixel_x, pixel_y
        
        pixel_point = np.array([[pixel_x, pixel_y]], dtype=np.float32)
        world_point = cv2.perspectiveTransform(
            pixel_point.reshape(-1, 1, 2), self.H
        ).reshape(-1, 2)[0]
        
        return world_point[0], world_point[1]
    
    def world_to_pixel(self, world_x, world_y):
        """Convert world coordinates to pixel coordinates"""
        if self.H is None:
            return world_x, world_y
        
        # 计算单应性矩阵的逆
        H_inv = np.linalg.inv(self.H)
        world_point = np.array([[world_x, world_y]], dtype=np.float32)
        pixel_point = cv2.perspectiveTransform(
            world_point.reshape(-1, 1, 2), H_inv
        ).reshape(-1, 2)[0]
        
        return pixel_point[0], pixel_point[1]
    
    def calculate_movement_to_center(self, current_world):
        """计算移动到图像中心所需的电机移动距离"""
        if current_world is None:
            return (0, 0)
        
        # 将图像中心坐标(325, 235)转换为世界坐标
        center_world_x, center_world_y = self.pixel_to_world(325, 235)
        
        # 计算需要移动的距离（X轴取反）
        dx = -(center_world_x - current_world[0])  # X轴距离取反
        dy = center_world_y - current_world[1]
        
        return (dx, dy)
    
    def is_at_center(self, pixel_x, pixel_y):
        """检查圆是否已经在中心位置"""
        distance = np.sqrt((pixel_x - 325)**2 + (pixel_y - 235)**2)
        return distance <= self.tolerance_pixels
    
    def auto_track_to_center(self, current_pixel):
        """自动追踪到中心位置"""
        if not self.serial_connected:
            print("串口未连接，无法进行自动追踪")
            return False
        
        if current_pixel is None:
            print("未检测到圆，无法追踪")
            return False
        
        # 检查是否已经在中心
        if self.is_at_center(current_pixel[0], current_pixel[1]):
            print("圆已经在中心位置，无需移动")
            return True
        
        # 计算世界坐标和移动距离
        world_x, world_y = self.pixel_to_world(current_pixel[0], current_pixel[1])
        movement = self.calculate_movement_to_center((world_x, world_y))
        
        # 检查移动距离是否过小
        if abs(movement[0]) < self.min_movement_threshold and abs(movement[1]) < self.min_movement_threshold:
            print(f"移动距离过小(X{movement[0]:.3f}, Y{movement[1]:.3f})，跳过移动")
            return True
        
        print(f"自动追踪: 需要移动 X{movement[0]:.2f}mm, Y{movement[1]:.2f}mm")
        
        # 执行移动
        success = self.move_motor_to_center(movement[0], movement[1])
        
        if success:
            print("移动完成，继续检测...")
            # 等待电机稳定
            time.sleep(0.5)
        else:
            print("移动失败")
        
        return success
    
    def start_auto_tracking(self):
        """开始自动追踪"""
        if not self.serial_connected:
            print("串口未连接，无法开始自动追踪")
            return
        
        self.tracking_enabled = True
        print("开始自动追踪...")
        print("将自动调整电机位置，直到黑圆移动到图像中心")
        print("按 's' 键停止自动追踪")
    
    def stop_auto_tracking(self):
        """停止自动追踪"""
        self.tracking_enabled = False
        print("自动追踪已停止")
    
    def manual_motor_control(self):
        """手动控制电机界面"""
        if not self.serial_connected:
            print("串口未连接，无法手动控制电机")
            return
            
        print("\n=== 手动电机控制 ===")
        print("输入移动距离 (格式: X10.5 Y-5.2), 输入 'q' 退出")
        
        while True:
            try:
                user_input = input("请输入移动距离: ").strip()
                if user_input.lower() == 'q':
                    break
                
                # 解析输入
                parts = user_input.split()
                x_move = 0
                y_move = 0
                
                for part in parts:
                    if part.lower().startswith('x'):
                        x_move = float(part[1:])
                    elif part.lower().startswith('y'):
                        y_move = float(part[1:])
                
                if x_move != 0 or y_move != 0:
                    success = self.move_motor_to_center(x_move, y_move)
                    if success:
                        print("移动完成")
                    else:
                        print("移动失败")
                else:
                    print("未输入有效移动距离")
                    
            except ValueError:
                print("输入格式错误，请使用格式: X10.5 Y-5.2")
            except Exception as e:
                print(f"错误: {e}")
    
    def run_detection(self):
        """Run black circle detection and validation program"""
        print("=== 黑圆检测和动态追踪程序 ===")
        print("按键控制:")
        print("  SPACE - 单次移动到图像中心")
        print("  'a'   - 开始/停止自动追踪")
        print("  'm'   - 手动控制电机移动")
        print("  'q'   - 退出程序")
        
        if not self.serial_connected:
            print("警告: 串口未连接，电机控制功能不可用")
            print("将只显示检测结果，无法控制电机")
        
        while True:
            ret, frame = self.cap.read()
            if not ret:
                print("Failed to read frame")
                break
            
            # Undistort image
            undistorted_frame = self.undistort_image(frame)
            
            # Detect black circles
            circles, binary = self.detect_black_circles(undistorted_frame)
            
            # Process detection results
            display_frame = undistorted_frame.copy()
            current_circle = None
            current_world = None
            
            # 绘制图像中心十字线和容差范围
            cv2.line(display_frame, (325, 225), (325, 245), (0, 0, 255), 2)
            cv2.line(display_frame, (315, 235), (335, 235), (0, 0, 255), 2)
            cv2.circle(display_frame, (325, 235), self.tolerance_pixels, (0, 255, 255), 1)
            cv2.putText(display_frame, "Center (325,235)", (250, 220), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            
            if circles:
                # Select the largest circle
                largest_circle = max(circles, key=lambda circle: circle[2])
                x, y, r = largest_circle
                current_circle = (x, y, r)
                
                # Convert to world coordinates
                world_x, world_y = self.pixel_to_world(x, y)
                current_world = (world_x, world_y)
                self.current_position = current_world
                
                # 检查是否在中心
                at_center = self.is_at_center(x, y)
                circle_color = (0, 255, 0) if at_center else (255, 0, 0)  # 绿色=在中心，蓝色=不在中心
                
                # Draw results on image
                cv2.circle(display_frame, (x, y), r, circle_color, 2)
                cv2.circle(display_frame, (x, y), 2, circle_color, 3)
                cv2.putText(display_frame, f"({x}, {y})", (x + 10, y - 10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                cv2.putText(display_frame, f"World: ({world_x:.1f}, {world_y:.1f})", 
                           (x + 10, y + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
                
                # 自动追踪逻辑
                if self.tracking_enabled and self.serial_connected:
                    if not at_center:
                        self.auto_track_to_center((x, y))
                    else:
                        # 在中心位置，显示成功信息
                        cv2.putText(display_frame, "TRACKING SUCCESS!", (200, 400), 
                                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
            
            # Display information
            cv2.putText(display_frame, f"Circles: {len(circles)}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
            
            # Display current movement info
            if current_world:
                movement = self.calculate_movement_to_center(current_world)
                move_text = f"To Center: X{movement[0]:+.1f}mm, Y{movement[1]:+.1f}mm"
                cv2.putText(display_frame, move_text, (10, 60), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            
            # 显示追踪状态
            tracking_status = "AUTO TRACKING: ON" if self.tracking_enabled else "AUTO TRACKING: OFF"
            tracking_color = (0, 255, 0) if self.tracking_enabled else (0, 0, 255)
            cv2.putText(display_frame, tracking_status, (10, 90), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, tracking_color, 2)
            
            # 显示串口状态
            status_text = "Motor: Connected" if self.serial_connected else "Motor: Disconnected"
            status_color = (0, 255, 0) if self.serial_connected else (0, 0, 255)
            cv2.putText(display_frame, status_text, (10, 120), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, status_color, 2)
            
            cv2.putText(display_frame, "SPACE:Single move, A:Auto track, M:Manual, Q:Quit", (10, 150), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            
            cv2.imshow('Black Circle Detection', display_frame)
            cv2.imshow('Binary View', binary)
            
            key = cv2.waitKey(1) & 0xFF
            if key == ord(' '):  # 单次移动到中心
                if not self.serial_connected:
                    print("串口未连接，无法控制电机")
                    continue
                    
                if current_world:
                    movement = self.calculate_movement_to_center(current_world)
                    print(f"\n=== 单次移动 ===")
                    print(f"需要移动的距离: X{movement[0]:.2f}mm, Y{movement[1]:.2f}mm")
                    
                    # 控制电机移动
                    success = self.move_motor_to_center(movement[0], movement[1])
                    if success:
                        print("移动完成")
                    else:
                        print("移动失败")
                else:
                    print("未检测到圆，无法移动")
            
            elif key == ord('a'):  # 开始/停止自动追踪
                if not self.serial_connected:
                    print("串口未连接，无法自动追踪")
                else:
                    if self.tracking_enabled:
                        self.stop_auto_tracking()
                    else:
                        self.start_auto_tracking()
                
            elif key == ord('m'):  # 手动控制电机
                if not self.serial_connected:
                    print("串口未连接，无法手动控制电机")
                else:
                    self.manual_motor_control()
                
            elif key == ord('q'):  # Quit
                break
        
        self.cap.release()
        cv2.destroyAllWindows()
        if self.ser:
            self.ser.close()
        print("Program exited")

# Usage example
if __name__ == "__main__":
    try:
        detector = CircleDetectionAndValidation(
            "hand_eye_calibration_results.txt", 
            "/dev/ttyUSB0"
        )
        detector.run_detection()
    except Exception as e:
        print(f"Program error: {e}")